tensor([[0.9700, 0.0800, 0.9700, 0.0800, 0.9700],
        [0.0300, 0.9200, 0.0300, 0.9200, 0.0300]], grad_fn=<RoundBackward1>)